import torch
import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # C1层：卷积层（输入通道1，输出通道6，卷积核5x5）
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        # S2层：下采样层（平均池化）
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
        # C3层：卷积层（输入通道6，输出通道16，卷积核5x5）
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        # S4层：下采样层（平均池化）
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
        # C5层：卷积层（输入通道16，输出通道120，卷积核5x5）
        self.conv3 = nn.Conv2d(16, 120, kernel_size=5)
        # 全连接层：F6层（输入120，输出84）
        self.fc1 = nn.Linear(120, 84)
        # 输出层：分类为10类
        # self.fc2 = nn.Linear(84, 10)

    def forward(self, x):
        # C1层 + ReLU + S2层
        x = self.pool1(F.relu(self.conv1(x)))
        # C3层 + ReLU + S4层
        x = self.pool2(F.relu(self.conv2(x)))
        # C5层 + ReLU
        x = F.relu(self.conv3(x))  # 输出维度为 [batch_size, 120, 1, 1]
        x = x.view(-1, 120)  # 展平为 [batch_size, 120]
        # F6层 + ReLU
        x = F.relu(self.fc1(x))
        # 输出层（分类结果）
        # x = self.fc2(x)
        return x


if __name__ == "__main__":
    model = LeNet()
    print(f'Total parameters: {sum(param.numel() for param in model.parameters())} ')
    # 遍历模型参数并输出每一层的名字和参数量
    for name, param in model.named_parameters():
        print(f"层名: {name}, 参数数量: {param.numel()}")
    x = torch.randn((1, 1, 32, 32))
    preds = model(x)
    print('输出shape:', preds.shape)
